#include <CEstimator.h>
#include <CCounter.h>


void CEstimator::Setup(string filename) {
  ini.load(filename);

  params.Setup(ini);


  cutoff_start_smallscale = ini.getdouble("cutoff_start_smallscale");

  cout << "Loading template " << ini.getstring("template_file") << endl;

  dataTemplate.LoadText( ini.getstring("template_file") );

  dataTemplate.FFT();
  dataTemplate.cutModes(cutoff_start_smallscale, dataTemplate.Bins);
  dataTemplate.FFTInv();

  dataTemplateOrg.Copy(dataTemplate);


  data.resize( (int)ini.getdouble("number_of_data_files"));
  basename = ini.getstring("base_data_name");

  for (int i=0;i<data.size(); i++) {
    string fname = basename + CUtil::toString(i+1) + ".txt";
    cout << "Loading file " << fname << endl;
    data[i].LoadText(fname);
  }

  cutoff_mode_start = ini.getint("cutoff_mode_start");
  cutoff_mode_end = ini.getint("cutoff_mode_end");

  cout << "Cutting modes " << cutoff_mode_start << " to " << cutoff_mode_end << endl;
  for (int i=0;i<data.size();i++) {
    data[i].FFT();
    data[i].cutModes(cutoff_mode_start, cutoff_mode_end);

    data[i].cutModes(cutoff_start_smallscale, data[i].Bins);


    data[i].FFTInv();
    //data[i].dampNoise(0.003, 0.0);
  }
  if (ini.getbool("output_washed_data")) {
    string base = ini.getstring("washed_data_base_name");
    for (int i=0;i<data.size();i++) {
      string fname = base + CUtil::toString(i+1) + ".txt";
      cout << "Saving file " << fname << endl;
      data[i].SaveText(fname);
      
    }
  }

  //cout << endl << "z = [ " << params.zfrom << " to " << params.zto << " ] with " << params.zsteps << " steps " << endl;

  cout << "loading absorption lines" << endl;

  use_absorption_lines = ini.getbool("use_absorption_lines");
    if (use_absorption_lines)
      setupAbsorption();
}




void CEstimator::setupAbsorption() {
  absorptionLines.LoadText( ini.getstring("absorptionlines_file"));

  double absorptionSigma = ini.getdouble("absorptionlines_sigma");
  double amplitude = 1.0;

  dataAbsorption.Copy(dataTemplate);

  for (int i=0;i<dataAbsorption.Bins;i++)
    dataAbsorption.Val[i] = 0.0;

  for (int i=0;i<absorptionLines.Bins; i++) {
    dataAbsorption.RenderGauss(absorptionLines.Index[i], absorptionSigma, amplitude);
  }

  dataAbsorption.SaveText("abs.txt");

  //  dataTemplate.Mul(dataAbsorption);

  dataTemplate.SaveText("temp1.txt");

}



void CEstimator::CalculateChisq1D() {
 
  for (int i=0;i<data.size();i++) {
    
    dataTemplateOrg.Std();
    data[i].Std();
    dataTemplate.Copy(dataTemplateOrg);
    dataTemplate.Scale(data[i].std / dataTemplateOrg.std);
    dataTemplate.Std();
    
    if (use_absorption_lines) 
      dataTemplate.Mul(dataAbsorption);

    dataTemplate.SaveText("template_" + CUtil::toString(i+1) +".txt");

    cout << dataTemplate.std << " vs " << data[i].std << endl;

    CalculateChisq1D(data[i], params.zfrom+1, params.zto+1, params.zsteps);
    string fname = basename + "_chisq" + CUtil::toString(i+1) + ".txt";
    cout << "Saving chisq file " << fname << endl;

    /*    data[i].chisq->Std();
    data[i].chisq->Scale(1.0/data[i].chisq->std);
    data[i].chisq->SaveText(fname);
    */
    data[i].chisq->Std();
    data[i].chisq->Scale(1.0/data[i].chisq->std);
    data[i].chisq->LikelihoodFromChisq();
    
    data[i].chisq->calculateStatistics();

    cout << "****** STATISTICS " << endl;
    cout << "Mean value : " << data[i].chisq->mean << " with std " << data[i].chisq->std << endl;



    data[i].chisq->SaveText(fname);
    

  }
}


void CEstimator::CalculateChisq2D() {
 
  for (int i=0;i<data.size();i++) {
    
    dataTemplateOrg.Std();
    data[i].Std();
    dataTemplate.Copy(dataTemplateOrg);
    dataTemplate.Scale(data[i].std / dataTemplateOrg.std);
    dataTemplate.Std();
    
    if (use_absorption_lines) 
      dataTemplate.Mul(dataAbsorption);
    
    dataTemplate.SaveText("template_" + CUtil::toString(i+1) +".txt");

    cout << dataTemplate.std << " vs " << data[i].std << endl;

    CalculateChisq2D(ini.getstring("2d_likelihood_file"), data[i]);
    /*string fname = basename + "_chisq" + CUtil::toString(i+1) + ".txt";
    cout << "Saving chisq file " << fname << endl;

        data[i].chisq->Std();
    data[i].chisq->Scale(1.0/data[i].chisq->std);
    data[i].chisq->LikelihoodFromChisq();
    
    data[i].chisq->calculateStatistics();

    cout << "****** STATISTICS " << endl;
    cout << "Mean value : " << data[i].chisq->mean << " with std " << data[i].chisq->std << endl;



    data[i].chisq->SaveText(fname);
    */

  }
}





void CEstimator::CalculateSmooth() {
  
  double s = ini.getdouble("smooth_length");
  string base = ini.getstring("washed_data_base_name");

  for (int i=0;i<data.size();i++) {
    data[i].Smooth(s);
    for (int i=0;i<data.size();i++) {
      string fname = base + CUtil::toString(i+1) + ".txt";
      cout << "Saving file " << fname << endl;
      data[i].SaveText(fname);
      
    }
  }
}

void CEstimator::Run() {

  if (ini.getstring("method")=="calculate_chisq_independently") {
      if (ini.getstring("bruteforce_likelihood_type")=="1D")
	CalculateChisq1D();
      if (ini.getstring("bruteforce_likelihood_type")=="2D")
	CalculateChisq2D();
    
  }

  if (ini.getstring("method")=="smooth_test") {
    CalculateSmooth();
    
  }

}


void CEstimator::CalculateChisq1D(CGraph& g, double zfrom, double zto, double steps) {
    double z = zfrom;
    double dz = (zto - zfrom)/steps;

    CGraph tmpGraph;

    g.chisq = new CGraph();
    g.chisq->Initialize((int)steps);
    for (int i=0;i<steps;i++) {
      //cout << z << endl;
      g.scaleX(1.0/z);
      
      tmpGraph.Copy(g);
      if (use_absorption_lines)
	tmpGraph.MulWithIndex(dataAbsorption);
	//tmpGraph.Mul(dataAbsorption);

      if (i==0) {
	  tmpGraph.SaveText("temp2.txt");
	  dataTemplate.SaveText("temp3.txt");

      }
      
      g.chisq->Val[i] = CGraph::ChiSQ(dataTemplate, tmpGraph);
      g.chisq->Index[i] = z - 1.0;
      g.chisq->IndexScaled[i] = z - 1.0;
      z += dz;
      //     cout << z << " " <<g.chisq->Val[i] << endl;

    }
}


void CEstimator::CalculateChisq2D(string file, CGraph& g) {

  CGraph tmpGraph;
  CGraph temp;
  
  ofstream f(file.c_str(), ios::out);
  CLikelihood like;
  like.Setup(params.zsteps, params.VDsteps);



  params.z = params.zfrom;

  vector<CGraph> vdsmooth;
  vdsmooth.resize(params.VDsteps);
  CCounter cntVD(params.VDsteps,"Setting up VD steps");
  for (int i=0;i<params.VDsteps;i++) {
    cntVD.Tick();
    vdsmooth[i].Copy(dataTemplate);
    vdsmooth[i].Smooth(params.VDfrom + params.dVD*i);
  }
  
  CCounter cnt(like.likelihood.size(),"Calculating 2D likelihood");
  for (int i=0;i<params.zsteps;i++) {

    g.scaleX(1.0/(params.z+1));
      
    tmpGraph.Copy(g);
    if (use_absorption_lines)
      tmpGraph.MulWithIndex(dataAbsorption);
    
    params.VD = params.VDfrom;
    for (int j=0;j<params.VDsteps;j++) {
      
      params.likelihood = CGraph::ChiSQ(vdsmooth[j], tmpGraph);

      like.likelihood[i + j*like.binx] = params;
      cnt.Tick();
      
      params.VD+=params.dVD;
    }
    params.z +=params.dz;
  }
  cout << endl;
  like.ToLikelihood();
  like.SaveText(file);
  like.Px.SaveText("px.txt");
  like.Py.SaveText("py.txt");

}


